"""
机器学习预测器模块
用于加载和使用训练好的MPI通信延迟预测模型
"""

import pickle
import numpy as np
import pandas as pd
import os
import sys
import warnings

class MPILatencyPredictor:
    """MPI通信延迟预测器"""
    
    def __init__(self, model_path=None):
        """
        初始化预测器
        
        Args:
            model_path: 模型文件路径，如果为None则使用默认路径
        """
        self.model = None
        self.model_path = model_path or self._get_default_model_path()
        self.load_model()
    
    def _get_default_model_path(self):
        """获取默认模型路径"""
        # 获取当前文件的目录
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 构建到MachineLearning_common/models的路径
        model_path = os.path.join(
            current_dir, 
            "..", "..", "..", "..", 
            "MachineLearning_common", 
            "models", 
            "mpi_latency_model.pkl"
        )
        return os.path.normpath(model_path)
    
    def load_model(self):
        """加载训练好的模型"""
        try:
            if os.path.exists(self.model_path):
                # 尝试使用joblib加载（推荐方式）
                try:
                    import joblib
                    model_data = joblib.load(self.model_path)
                except ImportError:
                    # 如果joblib不可用，使用pickle
                    with open(self.model_path, 'rb') as f:
                        model_data = pickle.load(f)
                
                # 检查模型数据格式
                if isinstance(model_data, dict) and 'model' in model_data:
                    # 新格式：包含模型、特征名称等的字典
                    self.model = model_data['model']
                    self.feature_names = model_data.get('feature_names', ['comm_type', 'message_size'])
                    print(f"成功加载ML模型: {self.model_path}")
                    print(f"模型类型: {type(self.model)}")
                elif hasattr(model_data, 'predict'):
                    # 直接的模型对象
                    self.model = model_data
                    self.feature_names = ['comm_type', 'message_size']
                    print(f"成功加载ML模型: {self.model_path}")
                else:
                    # 不支持的格式
                    print(f"警告: 不支持的模型格式 {type(model_data)}")
                    print("将使用备用预测方法")
                    self.model = None
            else:
                print(f"警告: 模型文件不存在 {self.model_path}")
                print("将使用备用预测方法")
                self.model = None
        except Exception as e:
            print(f"加载模型时出错: {e}")
            print("将使用备用预测方法")
            self.model = None
    
    def predict_single(self, comm_type, message_size):
        """
        单次预测
        
        Args:
            comm_type: 通信类型 (1-4)
            message_size: 消息大小（字节）
            
        Returns:
            预测的通信延迟（微秒）
        """
        if self.model is None:
            # 备用预测方法：使用简单的线性模型
            return self._fallback_prediction(message_size)
        
        try:
            # 构造完整的特征向量
            # 特征顺序：['type', 'message_size', 'log_message_size', 'size_category_encoded']
            log_message_size = np.log(max(1, message_size))  # 避免log(0)
            
            # 根据消息大小确定size_category_encoded
            if message_size <= 1024:
                size_category_encoded = 0  # small
            elif message_size <= 65536:
                size_category_encoded = 1  # medium
            else:
                size_category_encoded = 2  # large
            
            # 使用pandas DataFrame保持特征名称，避免sklearn警告
            feature_names = getattr(self, 'feature_names', ['type', 'message_size', 'log_message_size', 'size_category_encoded'])
            X = pd.DataFrame([[comm_type, message_size, log_message_size, size_category_encoded]], 
                           columns=feature_names)
            
            # 抑制sklearn的特征名称警告
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                prediction = self.model.predict(X)[0]
            
            return max(0, prediction)  # 确保预测值非负
            
        except Exception as e:
            print(f"ML预测出错: {e}")
            return self._fallback_prediction(message_size)
    
    def predict_average_all_types(self, message_size):
        """
        预测所有4种通信类型的平均延迟
        
        Args:
            message_size: 消息大小（字节）
            
        Returns:
            4种通信类型的平均预测延迟（微秒）
        """
        predictions = []
        for comm_type in [1, 2, 3, 4]:
            pred = self.predict_single(comm_type, message_size)
            predictions.append(pred)
        
        return np.mean(predictions)
    
    def _fallback_prediction(self, message_size):
        """
        备用预测方法（当ML模型不可用时）
        使用简单的线性模型
        """
        # 使用原始的C-LOP参数作为备用
        O = 1.076268
        L = 0.328143 * 0.001
        return O + L * message_size

# 全局预测器实例
_global_predictor = None

def get_ml_predictor():
    """获取全局ML预测器实例"""
    global _global_predictor
    if _global_predictor is None:
        _global_predictor = MPILatencyPredictor()
    return _global_predictor

def predict_ml_comm_time(message_size, comm_type=None):
    """
    便捷的ML预测函数
    
    Args:
        message_size: 消息大小（字节）
        comm_type: 通信类型，如果为None则返回4种类型的平均值
        
    Returns:
        预测的通信延迟（微秒）
    """
    predictor = get_ml_predictor()
    
    if comm_type is not None:
        return predictor.predict_single(comm_type, message_size)
    else:
        return predictor.predict_average_all_types(message_size)